from openai import OpenAI
import openai
import inspect
import json
import pymysql
import pandas as pd
import tiktoken
import os
import matplotlib.pyplot as plt
import tempfile
from docx import Document
from IPython.display import display, Markdown

openai.api_key = "sk-6S0PtpNia71gjcfwSsDPsJ9mGqsVPr2XRQzAx1dHbJS7RW4t"
openai.api_base = "https://chatapi.littlewheat.com/v1"

client = OpenAI(api_key=openai.api_key, base_url=openai.api_base)


def auto_functions(function_list):
    """
    Chat模型的functions参数编写函数
    :param functions_list: 包含一个或者多个函数对象的列表；
    :return：满足Chat模型functions参数要求的functions对象
    """

    def functions_generate(function_list):
        # 创建空列表，用于保存每个函数的描述字典
        functions = []
        # 对每个外部函数进行循环
        for function in function_list:
            # 读取函数对象的函数说明
            function_description = inspect.getdoc(function)
            # 读取函数的函数名字符串
            function_name = function.__name__
            system_prompt = "以下是某函数的说明:%s" % function_description
            user_prompt = '根据这个函数的函数说明，请帮我创建一个JSON格式的字典，这个字典有如下5点要求：\
                                       1.字典总共有三个键值对；\
                                       2.第一个键值对的Key是字符串name，value是该函数的名字：%s，也是字符串；\
                                       3.第二个键值对的Key是字符串description，value是该函数的函数的功能说明，也是字符串；\
                                       4.第三个键值对的Key是字符串parameters，value是一个JSON Schema对象，用于说明该函数的参数输入规范。\
                                       5.输出结果必须是一个JSON格式的字典，只输出这个字典即可，前后不需要任何前后修饰或说明的语句' % function_name
            response = client.chat.completions.create(
                model="gpt-3.5-turbo",
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt}
                ]
            )
            json_function_description = json.loads(
                response.choices[0].message.content.replace("```", "").replace("json", ""))
            json_str = {"type": "function", "function": json_function_description}
            functions.append(json_str)

        return functions

    max_attempts = 5
    attempts = 0
    while attempts < max_attempts:
        try:
            functions = functions_generate(function_list)
            # 如果代码成功执行，跳出循环
            break
        except Exception as e:
            # 增加尝试次数
            attempts += 1
            print("发生错误:", e)
            if attempts == max_attempts:
                print("已达到最大尝试次数，程序终止。")
                raise  # 重新引发最后一个异常
            else:
                print("正在重新运行...")
    return functions


def run_conversation(messages, function_list=None, model="gpt-3.5-turbo"):
    """
    能够自动执行外部函数调用的对话模型
    :param messages: 必要参数，字典类型，输入到Chat模型的messages参数对象
    :param functions_list: 可选参数，默认为None，可以设置为包含全部外部函数的列表对象
    :param model: Chat模型，可选参数，默认模型为gpt-3.5-turbo
    :return：Chat模型输出结果
    """
    # 如果没有外部函数库，则执行普通的对话任务
    if function_list == None:
        response = client.chat.completions.create(messages=messages, model=model)
        response_message = response.choices[0].message
        final_response = response_message.content
    # 若存在外部函数库，则需要灵活选取外部函数并进行回答
    else:
        # 创建functions对象
        tools = auto_functions(function_list)

        # 创建外部函数库字典
        available_functions = {func.__name__: func for func in function_list}

        # 第一次调用大模型
        response = client.chat.completions.create(messages=messages,
                                                  model=model,
                                                  tools=tools,
                                                  tool_choice="auto")
        response_message = response.choices[0].message
        tool_calls = response_message.tool_calls

        if tool_calls:
            messages.append(response_message)
            for tool_call in tool_calls:
                function_name = tool_call.function.name
                function_tool_call = available_functions[function_name]
                function_args = json.loads(tool_call.function.arguments)

                function_resp = function_tool_call(**function_args)
                messages.append({
                    "tool_call_id": tool_call.id,
                    "role": "tool",
                    "name": function_name,
                    "content": function_resp
                })

            ## 第二次调用模型
            response = client.chat.completions.create(
                model=model,
                messages=messages
            )

            # 获取最终结果
            final_response = response.choices[0].message.content
        else:
            final_response = response_message.content

        return final_response


def get_user_demographics(sql_query):
    """
    用户获取user_demographics 表的相关信息
    :param sql_query: 字符串形式的SQL语句
    :return SQL查询的user_demographics 表的相关信息
    """
    mysql_pw = "iquery_agent"
    connection = pymysql.connect(
        # 数据库地址
        host="localhost",
        # 数据库用户名
        user="iquery_agent",
        # 数据库密码
        passwd=mysql_pw,
        # 数据库名
        db="iquery",
        # 字符集选择utf8
        charset="utf8"
    )

    try:
        with connection.cursor() as cursor:
            sql = sql_query
            cursor.execute(sql)
            results = cursor.fetchall()
    finally:
        cursor.close()

    column_names = [desc[0] for desc in cursor.description]

    # 使用results和column_names创建DataFrame
    df = pd.DataFrame(results, columns=column_names)
    return df.to_json(orient="records")


with open('/root/autodl-tmp/chat_bi/数据字典/iquery数据字典.md', 'r', encoding='utf-8') as f:
    md_content = f.read()


def sql_inter(sql_query):
    """
    用于获取iquery数据库中各张表的有关相关信息，\
    核心功能是将输入的SQL代码传输至iquery数据库所在的MySQL环境中进行运行，\
    并最终返回SQL代码运行结果。需要注意的是，本函数是借助pymysql来连接MySQL数据库。
    :param sql_query: 字符串形式的SQL查询语句，用于执行对MySQL中iquery数据库中各张表进行查询，并获得各表中的各类相关信息
    :return：sql_query在MySQL中的运行结果。
    """
    mysql_pw = "iquery_agent"
    connection = pymysql.connect(
        # 数据库地址
        host="localhost",
        # 数据库用户名
        user="iquery_agent",
        # 数据库密码
        passwd=mysql_pw,
        # 数据库名
        db="iquery",
        # 字符集选择utf8
        charset="utf8"
    )

    try:
        with connection.cursor() as cursor:
            sql = sql_query
            cursor.execute(sql)
            results = cursor.fetchall()
    finally:
        cursor.close()

    return json.dumps(results)


# function_list = [sql_inter]
# tools = auto_functions(function_list)

def extract_sql(json_str):
    # 提取并返回'sql_query'的值
    return json_str.get('sql_query', None)


def check_code_run(messages,
                   project_name,
                   doc_name,
                   function_list=None,
                   tools=None,
                   model="gpt-4o",
                   auto_run=True):
    """
    能够自动执行外部函数调用的对话模型
    :param messages: 必要参数，字典类型，输入到Chat模型的messages参数对象
    :param functions_list: 可选参数，默认为None，可以设置为包含全部外部函数的列表对象
    :param model: Chat模型，可选参数，默认模型为gpt-3.5-turbo
    :return：Chat模型输出结果
    """
    # 如果没有外部函数库，则执行普通的对话任务
    if function_list == None:
        response = client.chat.completions.create(
            model=model,
            messages=messages
        )
        response_message = response.choices[0].message
        final_response = response_message.content
    # 若存在外部函数库，则需要灵活选取外部函数并进行回答
    else:
        # 创建functions对象
        tools = auto_functions(function_list)

        # 创建外部函数库字典
        available_function = {func.__name__: func for func in function_list}

        response = client.chat.completions.create(
            model=model,
            messages=messages,
            tools=tools,
            tool_choice="auto"
        )

        response_message = response.choices[0].message
        tool_calls = response_message.tool_calls
        if tool_calls:
            messages.append(response_message)

            for tool_call in tool_calls:
                function_name = tool_call.function.name
                function_tool_call = available_function[function_name]
                function_args = json.loads(tool_call.function.arguments)

                #########################################
                # 创建code对象
                if "sql_inter" in function_name or "extract_data" in function_name:
                    code = extract_sql(function_args)
                    # 将代码字符串转换为Markdown格式
                    markdown_code = f"```sql\n{code}\n```"
                else:
                    code = extract_python(function_args)
                    code = insert_fig_object(code)
                    # 将代码字符串转换为Markdown格式
                    markdown_code = f"```python\n{code}\n```"
                #########################################

                if auto_run == False:
                    sql_query = extract_sql(function_args)
                    res = input('即将执行以下代码：%s。是否确认并继续执行（1），或者退出本次运行过程（2）' % sql_query)
                    if res == '2':
                        print("终止运行")
                        return None
                    else:
                        print("正在执行代码，请稍后...")

                function_resp = function_tool_call(**function_args)
                messages.append({
                    "tool_call_id": tool_call.id,
                    "role": "tool",
                    "name": function_name,
                    "content": function_resp
                })

            ## 第二次调用模型
            response = client.chat.completions.create(
                model=model,
                messages=messages
            )
            final_response = response.choices[0].message.content
        else:
            final_response = response_message.content
    del messages
    ##################################
    if code:
        substrings = ['plt.', 'matplotlib.pyplot.', 'sns.', 'seaborn.', 'plot']
        if any(substring in code for substring in substrings):
            global fig
            append_img_in_doc(project_name, doc_name, img=fig)
            print("图片保存成功！")
    ##################################
    return final_response


# messages = [
#     {"role": "system", "content": md_content},
#     {"role": "user", "content": "请问iquery数据库下user_demographics表的第10条数据内容是？"}
# ]
# result = check_code_run(messages,
#                function_list=function_list,
#                model="gpt-3.5-turbo",
#                auto_run = False)
# print(result)

def chat_with_inter(function_list=None,
                    prompt="你好",
                    model="gpt-4o",
                    system_message=[{"role": "system", "content": "你是一个智能助手。"}],
                    auto_run=True
                    ):
    # 创建函数列表对应的参数解释列表
    tools = auto_functions(function_list)

    project_name = input("请输入当前分析项目名称：")
    folder_name = create_directory(project_name)
    print("已完成数据分析文件创建")
    doc_name = input("请输入当前分析需求，如数据清理，数据处理，数据分析段等：")
    doc_name += '问答'
    print("好的，即将进入交互式分析流程")

    # 多轮对话阈值
    if 'gpt-4' in model:
        tokens_thr = 6000
    elif '16k' in model:
        tokens_thr = 14000
    else:
        tokens_thr = 3000

    messages = system_message
    ## 完成给用户输入的问题赋值
    user_input = prompt
    messages.append({"role": "user", "content": user_input})
    ## 计算token大小
    embedding_model = "text-embedding-ada-002"
    # 模型对应的分词器（TOKENIZER）
    embedding_encoding = "cl100k_base"
    encoding = tiktoken.get_encoding(embedding_encoding)
    tokens_count = len(encoding.encode((prompt + system_message[0]["content"])))

    while True:
        answer = check_code_run(messages,
                                project_name=project_name,
                                doc_name=doc_name,
                                function_list=function_list,
                                tools=tools,
                                model=model,
                                auto_run=auto_run)
        print(f"模型回答: {answer}")

        #####################判断是否记录文档 start#######################
        while True:
            record = input('是否记录本次回答（1），还是再次输入问题并生成该问题答案（2）')
            if record == "1":
                q_tmp = "Q:" + user_input
                a_tmp = "A:" + answer
                append_in_doc(folder_name=project_name,
                              doc_name=doc_name,
                              qa_string=q_tmp)
                append_in_doc(folder_name=project_name,
                              doc_name=doc_name,
                              qa_string=a_tmp)
                # 记录本轮问题答案
                messages.append({"role": "assistant", "content": answer})
                break
            else:
                print('好的，请再次输入问题')
                user_input = input()
                messages[-1]["content"] = user_input
                answer = check_code_run(messages, function_list=function_list, model=model, auto_run=auto_run)
                print(f"模型回答: {answer}")
        ########################判断是否记录文档 stop #######################

        # 询问用户是否还有其他问题
        user_input = input("您还有其他问题吗？(输入退出以结束对话): ")
        if user_input == "退出":
            del messages
            break

        # 记录新一轮问答
        messages.append({"role": "assistant", "content": answer})
        messages.append({"role": "user", "content": user_input})

        # 计算当前总token数
        tokens_count += len(encoding.encode((answer + user_input)))

        # 删除超出token阈值的对话内容
        while tokens_count >= tokens_thr:
            tokens_count -= len(encoding.encode(messages.pop(1)["content"]))


def create_directory(directory):
    """
    根据项目创建云盘目录
    """
    base_path = "/root/autodl-tmp/chat_bi"
    full_path = os.path.join(base_path, directory)
    # 如果目录不存在，则创建它
    if not os.path.exists(full_path):
        os.makedirs(full_path)
        print(f"目录 {directory} 创建成功")
    else:
        print(f"目录 {directory} 已存在")


def append_in_doc(folder_name, doc_name, qa_string):
    """"
    往文件里追加内容
    @param folder_name=目录名，doc_name=文件名，qa_string=追加的内容
    """
    base_path = "/root/autodl-tmp/chat_bi"
    ## 目录地址
    full_path_folder = base_path + "/" + folder_name
    ## 文件地址
    full_path_doc = os.path.join(full_path_folder, doc_name) + ".doc"

    # 检查目录是否存在，如果不存在则创建
    if not os.path.exists(full_path_folder):
        os.makedirs(full_path_folder)

    # 检查文件是否存在
    if os.path.exists(full_path_doc):
        document = Document(full_path_doc)
    else:
        # 文件不存在，创建一个新的文档对象
        document = Document()

    # 追加内容
    document.add_paragraph(qa_string)
    # 保存文档
    document.save(full_path_doc)
    print(f"内容已追加到 {doc_name}")


# result = chat_with_inter(function_list=function_list,
#                 prompt="我想根据iquery数据库中数据分析用户流失情况，请问需要用到iquery数据库中的哪几张表呢？",
#                 model="gpt-3.5-turbo",
#                 system_message=[{"role": "system", "content": md_content}],
#                 auto_run = False)
# print(result)

# 读取业务知识文档
with open('/root/autodl-tmp/chat_bi/业务知识/本公司数据分析师业务介绍.md', 'r', encoding='utf-8') as f:
    da_instruct = f.read()

display(Markdown(da_instruct))


def get_file_content(project_name, file_name):
    """
        实现根据项目名和文件名获取文件内容的方法
        @param project_name:项目名，file_name：文件名
        @return 文件内容
        """
    # 构建文件的完整路径
    base_path = "/root/autodl-tmp/iquery项目/iquery云盘"
    file_path = os.path.join(project_name, file_name)
    full_path = os.path.join(base_path, file_path) + ".doc"
    print("打印文件路径：" + full_path)

    if not os.path.exists(full_path):
        return "文件不存在"

    try:
        # 加载文档
        doc = Document(full_path)
        content = []

        # 遍历文档中的每个段落，并收集文本
        for para in doc.paragraphs:
            content.append(para.text)

        # 将所有段落文本合并成一个字符串返回
        return "\n".join(content)
    except Exception as e:
        return f"读取文件时发生错误: {e}"


def extract_data(sql_query, df_name):
    """
    用于借助pymysql，将MySQL中的iquery数据库中的表读取并保存到本地Python环境中。
    :param sql_query: 字符串形式的SQL查询语句，用于提取MySQL中iquery数据库中的某张表。
    :param df_name: 将MySQL数据库中提取的表格进行本地保存时的变量名，以字符串形式表示。
    :return：表格读取和保存结果
    """
    mysql_pw = "iquery_agent"

    connection = pymysql.connect(
        host="localhost",
        user="iquery_agent",
        passwd=mysql_pw,
        db="iquery",
        charset="utf8"
    )

    globals()[df_name] = pd.read_sql(sql_query, connection)
    return "已成功完成%s变量创建" % df_name


sql_query = "SELECT * FROM user_demographics LIMIT 10"

data_name = "test_dataframe"
result = extract_data(sql_query, data_name)
print("test_dataframe:" + result)
print(test_dataframe)


# function_list = [extract_data]
# tools = auto_functions(function_list)

# messages=[
#             {"role": "user", "content": "请帮我读取iquery数据库中的user_payments，并保存至test_dataframe"}
#         ]

# response = client.chat.completions.create(
#         model="gpt-4o",
#         messages=messages,
#         tools=tools,
#         tool_choice="auto",
#     )
# print(response)

# messages = [
#     {"role": "user", "content": "读取iquery数据库中的user_payments数据表，并保存为test_dataframe"}
# ]

# response = client.chat.completions.create(
#     model="gpt-3.5-turbo",
#     messages=messages,
#     tools=tools,
#     tool_choice="auto",
# )
# print(response.choices[0].message)


# 能够在本地执行python代码的外部函数
# 定义一个名为python_inter的函数，该函数接受一个字符串类型的python代码作为参数
def python_inter(py_code):
    """
    该函数的主要作用是对iquery数据库中各张数据表进行查询和处理，并获取最终查询或处理结果。
    :param py_code: 字符串形式的Python代码，此代码用于执行对iquery数据库中各张数据表进行操作
    :return：返回代码运行的最终结果
    """
    # 添加图片对象，如果存在绘图代码，则创建fig对象
    py_code = insert_fig_object(py_code)
    # 记录函数开始执行时，全局作用域内的变量名
    global_vars_before = set(globals().keys())
    try:
        # 尝试执行传入的代码，代码的作用域为全局作用域
        exec(py_code, globals())
    except Exception as e:
        # 如果代码执行出错，返回错误信息
        return str(e)

    # 记录代码执行后，全局作用域内的变量名
    global_vars_after = set(globals().keys())
    # 获取执行代码后新产生的全局变量名
    new_vars = global_vars_before - global_vars_after
    # 如果有新的全局变量产生
    if new_vars:
        # 返回新的全局变量及其值
        result = {var: globals()[var] for var in new_vars}
        return str(result)
    else:
        try:
            print(test_dataframe)
            # 如果没有新的全局变量产生，尝试计算并返回代码的执行结果
            return str(exec(py_code, globals()))
        except Exception as e:
            # 如果计算代码执行结果也出错，返回成功执行代码的消息
            return "已经顺利执行代码"


functions_list = [python_inter]
tools = auto_functions(functions_list)

# messages=[
#             {"role": "user", "content": "已经读取iquery数据库中的user_payments数据表，并保存为test_dataframe"},
#             {"role": "user", "content": "test_dataframe是已经定义好的变量"},
#             {"role": "user", "content": "请帮我检查test_dataframe中各字段的基本类型和取值范围"}
#         ]

# response = client.chat.completions.create(
#         model="gpt-3.5-turbo-16k",
#         messages=messages,
#         tools=tools,
#         tool_choice="auto",
#     )

# messages = [
#     {"role": "user", "content": "已经读取iquery数据库中的user_payments数据表，并保存为test_dataframe"},
#     {"role": "user", "content": "test_dataframe是已经定义好的变量"},
#     {"role": "user", "content": "请帮我检查test_dataframe中总共包含几个字段"}
# ]

# messages = [
#     {"role": "user", "content": "已经读取iquery数据库中的user_payments数据表，并保存为test_dataframe"},
#     {"role": "user", "content": "test_dataframe是已经定义好的变量"},
#     {"role": "user", "content": "请帮我检查test_dataframe中各列的缺失值情况，并将统计结果保存为missing_value对象"}
# ]

messages = [
    {"role": "user", "content": "已经读取iquery数据库中的user_payments数据表，并保存为test_dataframe"},
    {"role": "user", "content": "test_dataframe是已经定义好的变量"},
    {"role": "user",
     "content": "请帮我统计test_dataframe中gender字段的不同取值分布情况，并以可视化形式进行展示，可视化展示时请尽量绘制更加清晰美观的图片"}
]

response = client.chat.completions.create(
    model="gpt-4o",
    messages=messages,
    tools=tools,
    tool_choice="auto",
)
print(response.choices[0].message)
py_code = response.choices[0].message.tool_calls[0].function.arguments
code = json.loads(py_code)
print(f"code:%s" % code)
code_str = code.get("py_code")
print("code_str:" + code_str)
result = python_inter(code_str)
print("result:" + result)
# print(missing_value)

# messages = [
#     {"role": "user", "content": "已经读取iquery数据库中的user_payments数据表，并保存为test_dataframe"},
#     {"role": "user", "content": "test_dataframe是已经定义好的变量"},
#     {"role": "user", "content": "请帮我检查test_dataframe中各列的缺失值情况，并将统计结果保存为missing_value对象"}
# ]

# response = client.chat.completions.create(
#     model="gpt-3.5-turbo",
#     messages=messages,
#     tools=tools,
#     tool_choice="auto",
# )

# message = response.choices[0].message
# print(message)

# code = json.loads(response.choices[0].message.tool_calls[0].function.arguments)
# code_str = code.get("py_code")
# result = python_inter(code_str)
# print(result)


def append_img_in_doc(folder_name,doc_name,img):
    fig,ax = plt.subplots()
    """"
    往文件里追加图片
    @param folder_name=目录名，doc_name=文件名，img=图片对象，数据类型为matplotlib.figure.Figure对象
    """
    base_path = "/root/autodl-tmp/chat_bi"
    ## 目录地址
    full_path_folder = base_path + "/" + folder_name
    ## 文件地址
    full_path_doc = os.path.join(full_path_folder,doc_name)+".doc"

    # 检查目录是否存在，如果不存在则创建
    if not os.path.exists(full_path_folder):
        os.makedirs(full_path_folder)

    # 检查文件是否存在
    if os.path.exists(full_path_doc):
        print(full_path_doc)
        # 文件存在，打开并追加内容
        document = Document(full_path_doc)
    else:
        # 文件不存在，创建一个新的文档对象
        document = Document()

    # 追加图片
    # 将matplotlib的Figure对象保存为临时图片文件
    with tempfile.NamedTemporaryFile(delete=False,suffix=".png") as tmpfile:
        fig.savefig(tmpfile.name,format="png")
        # 将图片插入到.docx文档中
        document.add_picture(tmpfile.name)

    # 保存文档
    document.save(full_path_doc)
    print(f"图片已追加到 {doc_name}")

def insert_fig_object(code_str):
    global fig
    # 检查是否已存在 fig 对象的创建
    if 'fig = plt.figure' in code_str:
        return code_str  # 如果存在，则返回原始代码字符串

    # 定义可能的库别名和全名
    plot_aliases = ['plt.', 'matplotlib.pyplot.', 'plot']
    sns_aliases = ['sns.', 'seaborn.']

    # 寻找第一次出现绘图相关代码的位置
    first_plot_occurrence = min((code_str.find(alias) for alias in plot_aliases + sns_aliases if code_str.find(alias) >= 0),default=-1)

    # 如果找到绘图代码，则在该位置之前插入 fig 对象的创建
    if first_plot_occurrence != "-1":
        plt_figure_index = code_str.find("plt.figure")
        if plt_figure_index != -1:
            # 寻找 plt.figure 后的括号位置，以确定是否有参数
            closing_bracket_index = code_str.find(')', plt_figure_index)
            # 如果找到了 plt.figure()，则替换为 fig = plt.figure()
            modified_str = code_str[:plt_figure_index] + 'fig = ' + code_str[plt_figure_index:closing_bracket_index + 1] + code_str[closing_bracket_index + 1:]
        else:
            modified_str = code_str[:first_plot_occurrence] + 'fig = plt.figure()\n' + code_str[first_plot_occurrence:]
        return modified_str
    else:
        return code_str  # 如果没有找到绘图代码，则返回原始代码字符串

def extract_python(json_str):
    if isinstance(json_str,dict):
        return json_str["py_code"]
    # 使用literal_eval将字符串转换为字典
    dict_data = json.loads(json_str)
    # 提取'sql_query'的值
    py_code_value = dict_data["py_code"]
    # 提取并返回'sql_query'的值
    return py_code_value

functions_list = [extract_data, python_inter]

chat_with_inter(functions_list=functions_list,
                prompt="请将iquery数据库中的user_demographics表保存到本地Python环境中，并命名为user_demographics_df",
                model="gpt-4o",
                system_message=[{"role": "system", "content": md_content}],
                auto_run = False)
